#!/usr/bin/env python3
# coding=utf-8

import rospy
import time
from wpr_yolo5.msg import BBox3D
from std_msgs.msg import String
from geometry_msgs.msg import Pose
from geometry_msgs.msg import Twist
from enum import Enum
from sensor_msgs.msg import JointState

# 抓取目标名称
grab_target = 'water'

# 任务状态
class state(Enum):
    Idle = 0
    GoToWP1      = 1
    ObjectDetect = 2
    GrabObject   = 3
    GoToWP2      = 4
    PutDown      = 5
    Backward     = 6
    Done         = 7

current_state = state.Idle

# 标记变量，是否处于抓取过程当中
grabbing = False

def cbNaviResult(msg):
    global current_state
    global behaviors_pub
    global mani_pub
    if current_state == state.GoToWP1:
        if msg.data == 'done':
            rospy.loginfo("到达取物地点，开始识别抓取")
            time.sleep(2)
            behavior_msg = String()
            behavior_msg.data = "yolo start"
            behaviors_pub.publish(behavior_msg)
            current_state = state.ObjectDetect

    if current_state == state.GoToWP2:
        if msg.data == 'done':
            rospy.loginfo("到达桌子，松手放置物品")
            mani_msg = JointState()
            mani_msg.name = ['gripper']
            mani_msg.position = [ 0.15 ]
            mani_msg.velocity = [ 5 ]
            mani_pub.publish(mani_msg)
            current_state = state.PutDown

# 物品检测回调函数
def cbObject(msg):
    global current_state
    global grabbing
    global behaviors_pub
    if current_state != state.ObjectDetect:
        return
    if grabbing == False:
        num = len(msg.name)
        rospy.loginfo("检测物品个数 = %d", num)
        if num > 0:
            i = 0
            for obj_name in msg.name:
                obj_x = msg.x_min[i]
                obj_y = (msg.y_min[i] + msg.y_max[i])/2
                obj_z = (msg.z_min[i] + msg.z_max[i])/2
                i = i+1
                rospy.loginfo("物品 %s! (%.2f , %.2f , %.2f)",obj_name,obj_x,obj_y,obj_z)
                if(obj_name == grab_target):
                    rospy.logwarn("抓取目标 %s! (%.2f , %.2f , %.2f)",obj_name,obj_x,obj_y,obj_z)
                    grab_msg = Pose()
                    grab_msg.position.x = obj_x
                    grab_msg.position.y = obj_y
                    grab_msg.position.z = obj_z
                    global grab_pub
                    grab_pub.publish(grab_msg)
                    grabbing = True
                    # 已经获取物品坐标，停止检测
                    behavior_msg = String()
                    behavior_msg.data = "yolo stop"
                    behaviors_pub.publish(behavior_msg)
                    current_state = state.GrabObject

# 抓取执行结果回调函数
def cbGrabResult(msg):
    global current_state
    global navi_pub
    if(current_state == state.GrabObject):
        if(msg.data == 'done'):
            rospy.loginfo("抓取完成，去往航点2")
            waypoint_msg = String()
            waypoint_msg.data = '2'
            navi_pub.publish(waypoint_msg)
            current_state = state.GoToWP2

# 主函数
if __name__ == "__main__":
    rospy.init_node("grab_node")
    # 发布导航目标点话题
    navi_pub = rospy.Publisher("/waterplus/navi_waypoint", String, queue_size=10)
    # 发布物品检测激活话题
    behaviors_pub = rospy.Publisher("/yolo_cmd", String, queue_size=10)
    # 发布抓取行为激活话题
    grab_pub = rospy.Publisher("/wpb_home/grab_action", Pose, queue_size=10)
    # 发布手臂控制指令的话题
    mani_pub = rospy.Publisher("/wpb_home/mani_ctrl", JointState, queue_size=10)
    # 发布速度控制指令的话题
    vel_pub = rospy.Publisher("/cmd_vel", Twist, queue_size=10)
    # 订阅导航结果的话题
    result_sub = rospy.Subscriber("/waterplus/navi_result", String, cbNaviResult, queue_size=10) 
    # 订阅物品检测结果的话题
    object_sub = rospy.Subscriber("/draw_bbox_3d", BBox3D, cbObject, queue_size=10)
    # 订阅抓取执行结果的话题
    result_sub = rospy.Subscriber("/wpb_home/grab_result", String, cbGrabResult, queue_size=10) 

    # 延时2秒，让后台的话题初始化操作能够完成
    rospy.sleep(2.0)

    count = 0

    rate = rospy.Rate(0.1)
    while not rospy.is_shutdown():
        if current_state == state.Idle:
            rospy.loginfo("开始去往取物地点")
            waypoint_msg = String()
            waypoint_msg.data = '1'
            navi_pub.publish(waypoint_msg)
            current_state = state.GoToWP1

        if current_state == state.PutDown:
            count += 1
            if count > 5*10:
                count = 0
                rospy.loginfo("物品已经松手，开始后退")
                current_state = state.Backward

        if current_state == state.Backward:
            count += 1
            vel_msg = Twist()
            if(count < 5*10):
                vel_msg.linear.x = -0.1
            else:
                vel_msg.linear.x = 0
                rospy.loginfo("任务完成，收回手臂")
                mani_msg = JointState()
                mani_msg.name = ['lift']
                mani_msg.position = [ 0 ]
                mani_msg.velocity = [ 0.5 ]
                mani_pub.publish(mani_msg)
                current_state = state.Done
            vel_pub.publish(vel_msg)

        rate.sleep()